package cn.xxi.ai.workflow.state;


import cn.xxi.ai.application.dto.ChatContext;
import cn.xxi.ai.application.dto.SseChatMessage;
import cn.xxi.ai.util.JsonUtil;
import dev.langchain4j.data.message.ChatMessage;
import org.bsc.langgraph4j.state.AgentState;
import reactor.core.publisher.Sinks;

import java.util.Map;

public class State extends AgentState {

    public State(Map<String, Object> initData) {
        super(initData);
    }

    @SuppressWarnings("unchecked")
    public <T> T get(String key) {
        return (T) data().get(key);
    }

    public <T> T get(ExtraKey key) {
        return get(key.name());
    }

    public Object any(String key) {
        return get(key);
    }

    public String string(String key) {
        Object value = get(key);
        return value != null ? value.toString() : null;
    }

    public String json() {
        return JsonUtil.toJson(data());
    }

    public ChatContext chatContext() {
        return get(ExtraKey.CHAT_CONTEXT);
    }

    public Sinks.Many<SseChatMessage> sinks() {
        return chatContext().sinks();
    }
}